import os
from pprint import pprint

import pandas as pd
from tabulate import tabulate

root = "experiments/test"
num_metrics = 4
num_dataset = 4

datasets = ["ML-1M", "ML-20M", "Beauty", "Game"]
metrics = ["Recall@5", "Recall@10", "NDCG@5", "NDCG@10"]
models = ["marank", "sas", "tisas", "bert", "meantime", "modify_meantime"]

Metrics = metrics * num_dataset

Dataset = []
for i in range(num_metrics * num_dataset):
    idx = i // num_metrics
    Dataset.append(datasets[idx])

print(Dataset)
print(Metrics)

data = {
    "Dataset": Dataset,
    "Metrics": Metrics,
    "MARANK": [],
    "SAS": [],
    "TISAS": [],
    "BERT": [],
    "MEANTIME": [],
    "MODIFY_MEANTIME": [],
}

dirs = os.listdir(root)
print(dirs)

dirs = sorted(dirs)
print(dirs)
# model_name = dir.split("_")[0]
# dataset_name = dir.split("_")[1]
# if model_name == "bert":
for model in models:
    for dataset in datasets:
        if dataset == "ML-1M":
            file_name = model + "_" + "1m"
        elif dataset == "ML-20M":
            file_name = model + "_" + "20m"
        elif dataset == "Beauty":
            file_name = model + "_" + "beauty"
        elif dataset == "Game":
            file_name = model + "_" + "game"
        else:
            pass

        df_path = os.path.join(root, file_name, "tables", "test_log.csv")
        df = pd.read_csv(df_path)

        recall_5 = "%.4f" % df["Recall@5"].iloc[0]
        recall_10 = "%.4f" % df["Recall@10"].iloc[0]
        ndcg_5 = "%.4f" % df["NDCG@5"].iloc[0]
        ndcg_10 = "%.4f" % df["NDCG@10"].iloc[0]

        data[model.upper()].extend((recall_5, recall_10, ndcg_5, ndcg_10))


# 计算改进
def calculate_improvement(row):
    scores = (float(score) for score in row[2:])
    sorted_scores = sorted(scores, reverse=True)  # 取得每行从第三列开始的得分，并排序
    if sorted_scores[1] == 0:  # 防止除数为零
        return None
    improvement = (sorted_scores[0] - sorted_scores[1]) / sorted_scores[1]  # 计算最高得分相对第二高得分的提升比率

    return "%.2f%%" % (improvement * 100)


# 创建DataFrame
df = pd.DataFrame(data)
# 设置打印时不要省略任何列
pd.set_option('display.max_columns', None)

df['Improvement'] = df.apply(calculate_improvement, axis=1)

# 使用tabulate来打印表格
print(tabulate(df, headers='keys', tablefmt='pretty', showindex=False))

df.to_csv("res.csv", index=None)
